fix(genrm_compare): keep a completed cohort through a checkpoint - #4792

A cohort whose judge fails gets neither the hold nor the export, so the stall this PR fixes still happens when the judge fails.

judge_failsafe turns the cohort's JudgeError into an HTTP 200 masked row (reward: 0.0, mask_sample: true). The caller records that row as the episode's result, just as it records a reward. But _fail_verify_cohort delivers it without waiting for resume, and _completed_member_state exports {} for any cohort whose phase isn't completed.

For example, say a checkpoint is open and the cohort is still evaluating, so the GenRM export writes {} for every member.

1. The judge deadline expires, and every member gets the masked row at once.
2. Member A's caller records the row before its own export, but member B's reply is processed after that export, so B is exported at its verify boundary.
3. After a crash, B replays into a process that has no cohort. It waits cohort_collection_timeout_s, then gets a 503 instead of the masked row that A recorded.

The same happens when the failure lands before admission closes, because a failed cohort is also exported as empty.

Could the export include failed cohorts that have a group ID, along with failure and failure_kind? Then a replayed verification would raise the same JudgeError and return the same masked row. Judge failures could also wait for resume, the same way rewards now do. If that's out of scope here, this comment should say that only completed cohorts are covered, since the last sentence currently reads as if every recorded result is.

Fixed in c8eaf759c. A cohort the judge failed now exports the failure, so after a restore a replayed verification raises the same JudgeError and judge_failsafe returns the same masked row. The judge-failure paths in the evaluation (a JudgeError and the evaluation deadline) are also held while a checkpoint is open, as rewards are. Cohort failures that callers don't record as a row (collection timeout, superseded attempt, shutdown) still export nothing and replay as before.

Tests: test_a_member_that_has_not_recorded_its_judge_failure_gets_it_back_after_a_restore, and test_outcomes_are_held_while_a_checkpoint_is_open[judge_failure]. Each fails if its half of the change is removed. Through the HTTP control plane, the judge-failure case on main stalls for the full collection timeout. Here it returns the masked row at once, after one or two checkpoints, without calling the judge.

When a restored session is retired or deleted, this hook forgets that session's member. But the restored cohort, and the watermark the restore raised, stay until cohort_result_ttl_s or the count cap evicts them. Until then they block the rollout from restarting from input.

Two controller paths retire restored sessions this way:

• A failed restore. Coordination retires the next attempts on every participant (coordination.py), then starts the rollouts over from their inputs.
• A commit that leaves a restored episode out of its scope. It calls delete_restored, whose docstring says the controller "may still start the rollout over as that attempt".

In both cases the group attempt stays the same. The README says a rollout's retry counter doesn't advance _ng_group_attempt, and the design doc says the group re-forms with the same attempt. So each regenerated response reaches the restored completed cohort and gets 409 "already has a different response for rollout_index=N". That lasts for up to an hour with the code default, or six hours with the README's example config. On main, the fresh process holds no cohort, so the group re-forms and completes.

flowchart LR
    R[restore installs completed cohort G attempt 0] --> F{another participant fails install, or commit scope drops the episode}
    F --> X[retire_session_state forgets the member only]
    X --> S[rollout restarts from input: new response, same group attempt]
    S --> C[409: already has a different response]

A smaller case has the same fix. A member that was reported as a restart, or that fell outside the commit scope, is missing from the restored cohort. When it reruns, it gets 409 "already completed; rollout_index=N cannot be added". That needs a group that mixes restarts with continued members, so it's narrow.

Could the restored cohort track which restored sessions back each member? Then, when a member's last restored session is retired or deleted before it replays, drop that member. Once no restored member is left, drop the cohort and the watermark the restore installed, so the group can re-form as it does on main.

Fixed in c8eaf759c, as suggested. A restore records which restored sessions back each member and haven't replayed yet. When such a session is retired or deleted before its member replays, through a failed restore's retire or delete_restored for a commit that leaves the episode out of scope, the member is dropped. A restored cohort left with no member is dropped together with the watermark its restore installed, unless live traffic has used that watermark since. The group then re-forms with the same attempt, as on main. A member that replayed stays, since it continued from the cohort.

Tests: test_a_restored_cohort_nothing_continues_from_is_dropped[failed_restore|left_out_of_commit] regenerate both members afterwards and complete the group. test_a_restored_cohort_stays_once_a_member_has_replayed covers the other side.

The narrower case you mention, a restart rerunning into a cohort that still holds continued members, still gets a 409 for its slot, and I've left it.

With checkpointing on, this records a member for every verification, including cohorts without _ng_group_id, and nothing ever frees those entries.

• A successful cohort without a group ID is popped in _publish_verify_cohort_locked, not in _prune_terminal_cohorts. The filter there runs only when that prune call itself drops a cohort, so it never runs when all traffic is like this.
• A session that ends normally, through /verify or /close_session, goes through ResourcesParticipant.ended(), which doesn't call retire_session_state.

So a run whose verifications carry no group ID adds one entry per rollout for the life of the process. A reproduction with 20 such cohorts and max_terminal_cohorts=2 was left holding 40 entries. The entries are never used either, because _completed_member_state returns {} for a cohort without a group ID. NeMo RL always sends a group ID, but every per-rollout structure still needs an end of life.

Recording only members of cohorts that have a group ID fixes it. The two call sites in verify would pass cohort instead of prompt_key:

Fixed in c8eaf759c along with the switch to session IDs. Only members of cohorts with a group ID are recorded, so a cohort without one adds nothing (test_a_cohort_without_a_group_id_records_no_members). Retiring a session pops its entry, and aclose clears the map.

This rebuilds both member maps on every prune call that drops any cohort. Once the server is at max_terminal_cohorts, that happens once per completed cohort. With the defaults (4096 cohorts, 16 members each), each rebuild walks about 65,000 entries while holding _cohort_registry_lock, which measured roughly 10 to 40 ms.

The existing sort in this prune is the same order of cost, so this isn't urgent. Still, keeping each cohort's member keys on _CohortState, and deleting them wherever the cohort is popped, would make cleanup proportional to what was evicted. Those places are the TTL prune, the count prune, the publish pop for cohorts without a group ID, and aclose. It would also give the publish pop an end-of-life path, and it maps directly onto #3625's eviction points, so the rebase gets easier.

Done in f75230a36. Each cohort keeps its member sessions, and the restored ones not replayed yet, on _CohortState. _drop_cohort removes exactly those records wherever a cohort goes: the TTL prune, the count prune, a dropped restore and aclose. Pruning no longer rebuilds the map. A cohort without a group ID records no members, so the publish pop has nothing to clean. test_member_records_end_with_their_pruned_cohort fails if the cleanup is removed.

GenRM could key its members by the resources session ID it already receives. That would remove this accessor and one of GenRM's two maps, and it needs no move to the typed session interface.

• ResourcesCheckpointMiddleware binds sessions by the cookie's SESSION_ID_KEY, and the participant passes those same IDs to export_session_states, restore_session_states and retire_session_state. If verify reads that ID too, GenRM's keys match the participant's on every path.
• The session middleware mints that ID whether an agent's /run seeded with the untyped body or an environment server sent a typed seed, where it equals resources_session_id. GenRM keeps no seed state, so the default seed_session and close_resources_session already cover both.
• judge_failsafe and traced_verify_endpoint both keep verify's signature with functools.wraps, so FastAPI injects a request parameter, as it does for workplace_assistant.
• A restore installs the same session IDs, and the restored boundary replays with the same cookie. So one map from session ID to member works both before and after a replay. That covers the restore-then-checkpoint case the second commit handles with _session_members.

A sketch:

async def verify(self, request: Request, body: GenRMCompareVerifyRequest) -> GenRMCompareVerifyResponse:
    session_id = request.session.get(SESSION_ID_KEY)
    ...
    # on both remember paths, instead of current_episode_id():
    if cohort.group_id is not None and session_id is not None and self.checkpoint_participant is not None:
        self._session_members[session_id] = (cohort.key, rollout_index)

Export then looks members up by session ID alone, restore fills the same map, and retire_session_state pops one key. _episode_members, the current_episode_id() lookup, and this new public method on the shared participant all go away. The tests that call server.verify(body) directly would need a request carrying a session. The map still needs the cohort cleanup from the comment on the prune filter.

Reordering install to bind sessions before it calls restore_session_states would also fix the second commit's case. But it changes the contract for every exported resources server and overlaps the install changes in #3909 and #4778, so it should be its own PR if anyone wants it.

Done in c8eaf759c. verify now takes request: Request and reads SESSION_ID_KEY, as workplace_assistant does, and one session-keyed map serves both before and after a replay. _episode_members, the current_episode_id() lookup and ResourcesParticipant.session_episode() are gone, so the PR no longer changes nemo_gym/. The existing tests that call verify directly now pass a session-less request built from the scope the session middleware leaves (tests/support.py). The checkpoint tests bind sessions on the real ResourcesParticipant.

#4220, which #4778 is stacked on, changes this hook to export_session_states(self, session_ids, *, checkpoint_id: str), and ResourcesParticipant.export passes checkpoint_id= to it. The two PRs merge without a textual conflict. But once both land, this override raises TypeError on every commit, so no checkpoint can commit on a deployment that runs GenRM.

A trial merge of #4778 into this branch reproduced it, first in test_checkpoint.py's participant.export(None) call and then here. Whichever PR lands second needs to add the keyword argument to this override and update those test calls. A note in both PR descriptions would keep that from being missed.

The digest saved here is whatever _response_digest produces. #3625 changes that function from json.dumps(..., ensure_ascii=True) to orjson.dumps, which writes UTF-8.

For any response with non-ASCII text the two digests differ. "café", for example, hashes differently under each. A checkpoint written before #3625 and restored after it would then get 409 "already has a different response" on every exact replay of such a member, and the reverse direction fails the same way.

Could this state carry a digest scheme field? This model sets extra="forbid", so a binary that doesn't know the field would refuse the state rather than misread it. A binary that does know it could recompute or skip the digest comparison when the scheme differs. Landing #3625 first would also avoid it, because no GenRM checkpoint would ever hold digests in the old format.

Done in ac04de348. Saved members now carry response_digest_scheme (RESPONSE_DIGEST_SCHEME = "json-ascii-sha256", defined next to _response_digest, with a comment that a serialization change must change the name). It's a fixed Literal under extra="forbid", so a restore refuses state saved under a scheme it doesn't know instead of installing digests it can't compare (test_a_restore_refuses_digests_saved_under_another_scheme). #3625 has been idle with conflicts since Sep 29, so I didn't rebase onto it. When it rebases, it can name its orjson scheme and skip the digest comparison for members restored under the old one. I've left a note there.

This loop leaves little margin. With the hold removed from _publish_verify_cohort, the test still passes at 4 iterations and only starts failing at about 6. await_count > 0 shows the judge was called, not that publication reached the hold. So one extra await in the evaluation path, which #3625's rewrite adds, would let this test pass without checking the hold.

Waiting until the publish is actually inside wait_open makes it deterministic, for example:

async def until(condition, timeout: float = 2) -> None:
    async with asyncio.timeout(timeout):
        while not condition():
            await asyncio.sleep(0)


def count_wait_open(participant) -> dict[str, int]:
    counts = {"entered": 0}
    wait_open = participant.wait_open

    async def counting_wait_open() -> None:
        counts["entered"] += 1
        await wait_open()

    participant.wait_open = counting_wait_open
    return counts

Then the test can replace the loop with await until(lambda: counts["entered"] == 1).

Fixed in c8eaf759c, using your until / count_wait_open approach: the test waits until the outcome is actually inside wait_open, and is parametrized over rewards and judge failures. f75230a36 adds a test that reopens and recloses admission while the outcome waits for the cohort lock. It covers the re-check under the lock for both publication and failure; without either re-check, the matching case fails.

Most of these tests run against this copy of the participant's surface instead of ResourcesParticipant. Only the last test uses the real one, so the rest can't catch a change in how the real participant binds sessions or reports admission.

Mutation runs also show that the end-of-life paths of the new maps, and one atomicity check, are untested. Each of these changes leaves every test in this file passing:

• making retire_session_state a no-op;
• removing both filters in _prune_terminal_cohorts;
• removing aclose's clearing of _episode_members and _session_members;
• removing the re-check of accepting under cohort.lock;
• skipping _remember_member on the existing-member path;
• removing the check that restored sessions agree on a member.

Building the tests on participant_server would close most of these gaps. Drive admission with close_admission/open_admission, bind sessions with seeded, and add tests asserting that the maps are empty after retire, prune and aclose. A test through the HTTP middleware is also missing: rollout prefix, close-follows header, then prepare, commit, restore, resume and replay. A version of that test passes on this branch and can be shared.

Addressed in c8eaf759c and f75230a36. test_checkpoint.py now runs on the real ResourcesParticipant: sessions bind with seeded, admission goes through close_admission/open_admission, and restores go through install, retire and delete_restored. Each mutation from your list now fails a test:

• retire_session_state as a no-op;
• the prune cleanup removed;
• aclose not clearing members;
• the re-check of accepting under cohort.lock, for both rewards and judge failures;
• skipping the member record on the existing-member (replay) path;
• the check that restored sessions agree on a member.

It also adds a test through the HTTP middleware: seed, verify with the close-follows header, prepare, commit, restore, resume and replay. It runs after one and after two checkpoint-crash-restore cycles, for rewards and for judge failures. The same flow on main stalls for the full collection timeout.

[P2] Keep saved outcomes while open sessions still need them.

At ac04de34, with max_terminal_cohorts=1, finish two two-member groups, close member 0 in each, then export and restore the remaining sessions. Both outcomes are saved, but replaying the older group prunes its outcome and waits for the sibling that already finished. This is a reachable checkpoint: pruning runs before collection, so completing the second group can leave two terminal groups until the next verification.

With the collection timeout shortened to 0.05 seconds, replay returns HTTP 503 and the next export contains empty state for that session. This affects saved judge failures as well as rewards.

Please retain outcomes needed by open checkpoint sessions independently of terminal-cache eviction, with regression coverage for both count-based and time-based pruning. The count-based case above is reproduced; the time-based path calls the same unconditional cleanup.

Thanks, reproduced. Fixed in 4ab8322.

A finished cohort is now kept while a member session the checkpoint participant still holds open belongs to it. Such a cohort skips both the cohort_result_ttl_s expiry and the max_terminal_cohorts count, and its group keeps its attempt watermark. "Open" is the participant's own view (is_open), not GenRM's member list: a member that records its outcome closes its session without a retire, so the member list alone would have kept nearly every cohort. Once those sessions close or are retired, the cohort is an ordinary tombstone again and is pruned as before. This also covers the live case: a cohort with an unrecorded member can no longer be evicted before a checkpoint exports it.

test_pruning_keeps_an_outcome_an_open_session_still_needs runs your scenario (two finished groups, member 0 of each closed, export, restore, replay both) by count (max_terminal_cohorts=1) and by age (cohort_result_ttl_s=0.01), for rewards and judge failures. Each replay returns its own outcome at once without calling the judge. While the cohorts are kept, both groups keep their watermarks. After the sessions close, pruning removes them again. test_member_records_end_with_their_pruned_cohort now closes both sessions before expecting the prune. Each of these mutations fails a test: keeping nothing, keeping any cohort with a member record (open or not), and dropping the kept cohort's watermark.

[P3, nonblocking] Use distinct final rewards to check member identity.

The new lock-race and checkpoint-control tests address those earlier gaps. The member-identity assertion still needs distinct final rewards: this fixture produces 3.0 for both members, so the assertion cannot detect a sibling mix-up.

I changed export to save the other member's result, and all 22 checkpoint tests still passed. Please use different final rewards and verify that each member recovers its own outcome, including after a second checkpoint.

This is a regression-test coverage request; I have not found incorrect reward assignment in the unmodified implementation.

Done in 4ab8322. The identity tests and the HTTP control-plane test now use a judge under which member 0 wins in either comparison order, so the two members finish with different rewards (each test asserts they differ). test_a_restored_cohort_survives_another_checkpoint_before_its_members_replay also runs with both members still open at the checkpoint, and checks that each gets back its own reward after two checkpoint-crash-restore cycles. With export changed to save the sibling's result, 8 tests now fail, including the second-checkpoint cases.

This keyword no longer exists on main. #4997 renamed verify_mode to verify_step_type on ResourcesParticipant, so on the merge with the latest main every test in this file fails while building the participant (TypeError: ResourcesParticipant.__init__() got an unexpected keyword argument 'verify_mode'). Git reports no conflict, so only a rebase or CI on the merge commit shows it.

With this one rename on the merged tree, the GenRM suites and the checkpoint unit tests pass, so app.py needs nothing for #3909, #4997 or #4679. Could you rebase onto main and apply:

Rebased onto main and applied in 2c5a654. The GenRM suites pass on the rebased branch (277 passed).

This keep check now runs participant.is_open for every member session of every finished cohort, on every verification that carries a group ID, while _cohort_registry_lock is held. Closed sessions stay in member_sessions, because the participant calls no hook when a session ends normally, so the scan covers them too.

At the defaults (4096 finished cohorts of 16 members), a prune measured a median of 7.6 ms with checkpointing on, against 3.1 ms with it off. That cost is paid about 16 times per cohort, once per member's verification, where the rebuild this replaced ran only when a cohort was evicted.

Could the check move to the eviction candidates only? In the time-based pass, test _needed_by_open_session only for cohorts whose age has passed cohort_result_ttl_s. In the count pass, walk oldest first and skip a needed cohort until the excess is dropped. The active-group set for watermarks then needs only the kept cohorts. The cost becomes proportional to what is evicted or kept rather than to everything retained, and the existing keep tests cover the behavior.

Done in 93cdc01, as suggested: the age pass checks _needed_by_open_session only for cohorts past cohort_result_ttl_s, and the count pass walks oldest first, skipping a needed cohort until the excess is dropped. Every retained cohort keeps its group's watermark. At 4096 finished cohorts of 16 members, a prune with checkpointing on went from ~10.8 ms to ~2.1 ms, the same as with it off.

One small change in what is kept: a needed cohort among the newest now counts toward max_terminal_cohorts, so needed cohorts only go beyond the cap when they would otherwise be evicted. I updated the README to match. Disabling each of the three keep checks makes the existing tests fail.